import networkx as nx
from sklearn.cluster import SpectralClustering


class Cluster(object):
    def __init__(self, Graph):
        self.G = Graph

    def fit(self, method, n_clusters=8):
        self.matrix = nx.to_numpy_matrix(self.G)

        if method == 'spectral_cluster':
            sc = SpectralClustering(n_clusters, affinity='precomputed', n_init=100, assign_labels='discretize')
            sc.fit(self.matrix)

            partition = {}
            nodes = nx.to_dict_of_lists(self.G).keys()
            for index, value in enumerate(nodes):
                partition[value] = sc.labels_[index]
        else:
            partition = None

        return partition
